"""
Paddle Varbase
"""

import paddle
from paddle import Tensor

Tensor.tmp = Tensor.__getitem__
def __getitem__(self, idx):
    """
    getitem function
    """
    if isinstance(idx, paddle.Tensor) and len(idx.shape) == 1:
        out = paddle.gather(self, idx)
        return out
    elif isinstance(idx, paddle.Tensor) and str(idx.dtype) == "VarType.BOOL":
        idx = paddle.cast(idx, "int32")
        idx = paddle.nonzero(idx)
        out = paddle.gather_nd(self, idx)
        return out
    elif isinstance(idx, tuple):
        return self.__getitem__(idx)

        # TODO(syf): 出来为(slice(None, None, None), slice(None, None, None), 0)
    else:
        return self.tmp(idx)
Tensor.__getitem__ = __getitem__


Tensor.setitem_tmp = Tensor.__setitem__
def __setitem__(self, idx, value):
    """
    setitem function
    """
    if isinstance(idx, paddle.Tensor) and str(idx.dtype) == "VarType.BOOL":
        value_tensor = paddle.full(self.shape, value, self.dtype)
        paddle.assign(paddle.where(idx, value_tensor, self), self)
    else:
        return self.setitem_(idx, value)
Tensor.__setitem__ = __setitem__